AWS CDK v2.115.0 で Step Functions の Bedrock InvokeModel API 最適化統合がサポートされました

AWS re:Invent 2023

#Amazon Bedrock

#AWS CDK

#AWS Step Functions

#TypeScript

若槻龍太

2023.12.16

こんにちは、CX 事業本部 Delivery 部の若槻です。

先月に AWS re:Invent 2023 期間のアップデートの一つとして、AWS Step Functions での Amazon Bedrock API の最適化統合のサポートがありました。

そして本日リリースされた AWS CDK v2.115.0 で、この最適化統合が L2 Construct でもサポートされました。

stepfunctions-tasks: support for the Step Functions optimized integration for Bedrock InvokeModel API (#28276) (f3dafa4)

試してみた

CDK ライブラリのアップグレード

AWS CDK のモジュールを v2.115.0 以上にアップグレードします。

npm i aws-cdk@latest aws-cdk-lib@latest

基礎モデルへのアクセスの有効化

使用したい基礎モデルへのアクセスが有効化されていることを確認します。

まだの場合は下記を参考にアクセスをリクエストし、有効化します。

今回は米国西部リージョンで使用できる Amazon Titan Text G1 - Express を使用したいので、アクセスを有効化しておきます。

型定義の確認

今回のアップデートで使用可能になった aws_stepfunctions_tasks.BedrockInvokeModel コンストラクトクラスの型定義を確認します。

クリックして開く

import { Construct } from 'constructs';
import * as bedrock from '../../../aws-bedrock';
import * as iam from '../../../aws-iam';
import * as s3 from '../../../aws-s3';
import * as sfn from '../../../aws-stepfunctions';
/**
 * Location to retrieve the input data, prior to calling Bedrock InvokeModel.
 *
 * @see https://docs.aws.amazon.com/step-functions/latest/dg/connect-bedrock.html
 */
export interface BedrockInvokeModelInputProps {
    /**
     * S3 object to retrieve the input data from.
     *
     * If the S3 location is not set, then the Body must be set.
     *
     * @default Input data is retrieved from the `body` field
     */
    readonly s3Location?: s3.Location;
}
/**
 * Location where the Bedrock InvokeModel API response is written.
 *
 * @see https://docs.aws.amazon.com/step-functions/latest/dg/connect-bedrock.html
 */
export interface BedrockInvokeModelOutputProps {
    /**
     * S3 object where the Bedrock InvokeModel API response is written.
     *
     * If you specify this field, the API response body is replaced with
     * a reference to the Amazon S3 location of the original output.
     *
     * @default Response body is returned in the task result
     */
    readonly s3Location?: s3.Location;
}
/**
 * Properties for invoking a Bedrock Model
 */
export interface BedrockInvokeModelProps extends sfn.TaskStateBaseProps {
    /**
     * The Bedrock model that the task will invoke.
     *
     * @see https://docs.aws.amazon.com/bedrock/latest/userguide/api-methods-run.html
     */
    readonly model: bedrock.IModel;
    /**
     * The input data for the Bedrock model invocation.
     *
     * The inference parameters contained in the body depend on the Bedrock model being used.
     *
     * The body must be in the format specified in the `contentType` field.
     * For example, if the content type is `application/json`, the body must be
     * JSON formatted.
     *
     * The body must be up to 256 KB in size. For input data that exceeds 256 KB,
     * use `input` instead to retrieve the input data from S3.
     *
     * You must specify either the `body` or the `input` field, but not both.
     *
     * @see https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters.html
     *
     * @default Input data is retrieved from the location specified in the `input` field
     */
    readonly body?: sfn.TaskInput;
    /**
     * The desired MIME type of the inference body in the response.
     *
     * @see https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_InvokeModel.html
     * @default 'application/json'
     */
    readonly accept?: string;
    /**
     * The MIME type of the input data in the request.
     *
     * @see https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_InvokeModel.html
     * @default 'application/json'
     */
    readonly contentType?: string;
    /**
     * The source location to retrieve the input data from.
     *
     * @default Input data is retrieved from the `body` field
     */
    readonly input?: BedrockInvokeModelInputProps;
    /**
     * The destination location where the API response is written.
     *
     * If you specify this field, the API response body is replaced with a reference to the
     * output location.
     *
     * @default The API response body is returned in the result.
     */
    readonly output?: BedrockInvokeModelOutputProps;
}
/**
 * A Step Functions Task to invoke a model in Bedrock.
 *
 */
export declare class BedrockInvokeModel extends sfn.TaskStateBase {
    private readonly props;
    private static readonly SUPPORTED_INTEGRATION_PATTERNS;
    protected readonly taskMetrics: sfn.TaskMetricsConfig | undefined;
    protected readonly taskPolicies: iam.PolicyStatement[] | undefined;
    private readonly integrationPattern;
    constructor(scope: Construct, id: string, props: BedrockInvokeModelProps);
    private renderPolicyStatements;
    /**
     * Provides the Bedrock InvokeModel service integration task configuration
     *
     * @internal
     */
    protected _renderTask(): any;
}

CDK で構築する

aws_stepfunctions_tasks.BedrockInvokeModel を使用するステートマシンを AWS CDK で構築します。

まず、CDK スタックを米国西部リージョンにデプロイするように env を指定します。

import { App } from 'aws-cdk-lib';
import { CdkSampleStack } from '../lib/cdk-sample-stack';

const app = new App();

new CdkSampleStack(app, 'CdkSampleStack', { env: { region: 'us-east-1' } });

CDK スタックのコードです。

import {
  aws_bedrock,
  aws_stepfunctions,
  aws_stepfunctions_tasks,
  Stack,
  StackProps,
} from 'aws-cdk-lib';
import { Construct } from 'constructs';

export class CdkSampleStack extends Stack {
  constructor(scope: Construct, id: string, props: StackProps) {
    super(scope, id, props);

    // モデルインスタンスの作成
    const model = aws_bedrock.FoundationModel.fromFoundationModelId(
      this,
      'Model',
      aws_bedrock.FoundationModelIdentifier.AMAZON_TITAN_TEXT_G1_EXPRESS_V1
    );

    // モデルを呼び出すタスクの作成
    const promptModelTask = new aws_stepfunctions_tasks.BedrockInvokeModel(
      this,
      'PromptModelTask',
      {
        model, // Bedrock モデルの指定
        body: aws_stepfunctions.TaskInput.fromObject({
          /**
           * @see https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-titan-text.html 
          */
          inputText: 'Generate a list of five first names.',
          textGenerationConfig: {
            maxTokenCount: 100,
            temperature: 1, // ランダム性の度合い
          },
        }),
        resultSelector: {
          names: aws_stepfunctions.JsonPath.stringAt(
            '$.Body.results[0].outputText'
          ),
        },
      }
    );

    // ステートマシンの作成
    new aws_stepfunctions.StateMachine(this, 'StateMachine', {
      definitionBody:
        aws_stepfunctions.DefinitionBody.fromChainable(promptModelTask),
    });
  }
}

model パラメーターでは使用する Bedrock モデルを指定します。モデルインスタンスは前回の記事で紹介した FoundationModel.fromFoundationModelId で取得できます。

body パラメーターではモデルの推論パラメーターを指定します。パラメーターの仕様はモデルによって異なるので以下のドキュメントを参照してください。

CDK デプロイをすると、ステートマシンが作成されました。

State Machine 定義は下記のようになります。今回はプロンプトで Generate a list of five first names.（5 つの姓のリストを生成して。）と固定の質問をするステートマシンとなっています。

{
  "StartAt": "PromptModelTask",
  "States": {
    "PromptModelTask": {
      "End": true,
      "Type": "Task",
      "ResultSelector": {
        "names.$": "$.Body.results[0].outputText"
      },
      "Resource": "arn:aws:states:::bedrock:invokeModel",
      "Parameters": {
        "ModelId": "arn:aws:bedrock:us-east-1::foundation-model/amazon.titan-text-express-v1",
        "Body": {
          "inputText": "Generate a list of five first names.",
          "textGenerationConfig": {
            "maxTokenCount": 100,
            "temperature": 1
          }
        }
      }
    }
  }
}

動作確認

動作確認のためにステートマシンを実行してみます。

実行が成功しました。

タスクのアウトプットは下記のようになりました。

{
  "Body": {
    "inputTextTokenCount": 8,
    "results": [
      {
        "tokenCount": 29,
        "outputText": " Here is a list of five first names:\n\n1. Josh\n2. Jordan\n3. Joseph\n4. Alexander\n5. Ethan",
        "completionReason": "FINISH"
      }
    ]
  },
  "ContentType": "application/json"
}

次の 5 つの姓を生成してくれました。質問に対して適切な回答を得ることができました。

Josh
Jordan
Joseph
Alexander
Ethan

同じステートマシンを再度実行すると、当然ですが異なる回答が返ってきます。

{
  "Body": {
    "inputTextTokenCount": 8,
    "results": [
      {
        "tokenCount": 27,
        "outputText": "\nHere are five first names:\n\n1. Emma\n2. James\n3. Olivia\n4. Sophia\n5. William",
        "completionReason": "FINISH"
      }
    ]
  },
  "ContentType": "application/json"
}

アクセスが有効化されていないモデルを使用した場合

アクセスが有効化されていないモデルを使用した場合はどうなるでしょうか。

前述のステートマシンで、Amazon Titan Text G1 - Express へのアクセスが無効化された状態でステートマシンを実行してみます。

するとステートマシンの実行で想定通りエラーが発生しました。Bedrock.AccessDeniedExceptionが発生しています。

You don't have access to the model with the specified model ID. (Service: BedrockRuntime, Status Code: 403, Request ID: e63fc6c2-832f-4203-a12a-1e52c825cd4d)

アクセスが無効なモデルを使用する場合でも CDK デプロイ時にエラーで落とされる挙動とはならなかったので、使用するモデルが有効化されているかどうかは確認が必要です。